package animation;

import com.crunch.math.Vector2f;
import org.json.JSONArray;
import org.json.JSONException;
import org.json.JSONObject;
import org.json.JSONTokener;

import java.io.IOException;
import java.io.Reader;
import java.io.Writer;
import java.nio.charset.Charset;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;

/**
 * A skeleton used for character animation.
 */
public class Skeleton {
	// initializes empty skeleton
	public Skeleton() {
		bones = new ArrayList<Bone>();
		rootBones = new ArrayList<Integer>();
		boneNames = new HashMap<String, Integer>();
		boneDrawingOrder = new ArrayList<Integer>();
	}

	// saves skeleton to file
	public void save(Path path)
			throws IOException {
		Writer writer = null;
		try {
			JSONObject rootObj = new JSONObject();

			JSONArray bonesArr = new JSONArray("bones");
			rootObj.put("bones", bonesArr);
			saveBones(bonesArr, rootBones);

			writer = Files.newBufferedWriter(path, Charset.defaultCharset());
			rootObj.write(writer);
		} catch (JSONException e) {
			System.out.println("Error: " + e.getMessage());
			throw new IOException("Failed to write " + path);
		} finally {
			if (writer != null) {
				try {
					writer.close();
				} catch (IOException e) {}
			}
		}
	}

	private void saveBones(JSONArray bonesArr, ArrayList<Integer> boneIndices) {
		for (int i = 0; i < boneIndices.size(); ++i) {
			Bone bone = getBone(boneIndices.get(i));
			JSONObject boneObj = new JSONObject("bone");
			bonesArr.put(boneObj);
			boneObj.put("name", bone.getName());
			boneObj.put("length", bone.getLength());
			boneObj.put("thickness", bone.getThickness());
			Vector2f translation = bone.getTranslation();
			boneObj.put("translation", new JSONArray().put(translation.x()).put(translation.y()));
			boneObj.put("rotation", bone.getRotation());
			Vector2f scale = bone.getScale();
			boneObj.put("scale", new JSONArray().put(scale.x()).put(scale.y()));
			boneObj.put("depth", bone.getDepth());

			JSONArray childrenArr = new JSONArray("children");
			boneObj.put("children", childrenArr);
			saveBones(childrenArr, bone.children);
		}
	}

	// loads skeleton from file
	public void load(Path path)
			throws IOException {
		bones.clear();
		rootBones.clear();
		boneNames.clear();
		boneDrawingOrder.clear();
		Reader reader = null;
		try {
			reader = Files.newBufferedReader(path, Charset.defaultCharset());
			JSONObject rootObj = new JSONObject(new JSONTokener(reader));

			JSONArray bonesArr = rootObj.getJSONArray("bones");
			loadBones(bonesArr, -1);
		} catch (JSONException e) {
			System.out.println("Error: " + e.getMessage());
			throw new IOException("Failed to read " + path);
		} catch (IOException e) {
			System.out.println("Error: " + e.getMessage());
			throw new IOException("Failed to read " + path);
		} finally {
			if (reader != null) {
				try {
					reader.close();
				} catch (IOException e) {}
			}
		}
	}

	private void loadBones(JSONArray bonesArr, int parent)
			throws IOException {
		for (int b = 0; b < bonesArr.length(); ++b) {
			JSONObject boneObj = bonesArr.getJSONObject(b);

			int index = bones.size();

			String name = boneObj.getString("name");
			float length = (float) boneObj.getDouble("length");
			float thickness = (float) boneObj.getDouble("thickness");
			JSONArray translationArr = boneObj.getJSONArray("translation");
			Vector2f translation = new Vector2f(
					(float) translationArr.getDouble(0),
					(float) translationArr.getDouble(1));
			float rotation = (float) boneObj.getDouble("rotation");
			JSONArray scaleArr = boneObj.getJSONArray("scale");
			Vector2f scale = new Vector2f(
					(float) scaleArr.getDouble(0),
					(float) scaleArr.getDouble(1));
			int depth = boneObj.getInt("depth");

			if (boneNames.containsKey(name)) {
				throw new IOException("Bone " + name + " already exists");
			}

			if (parent < 0) {
				addBone(name, length, thickness, translation, rotation, scale, depth);
			} else {
				addBone(name, parent, length, thickness, translation, rotation, scale, depth);
			}

			// add children
			JSONArray childrenArr = boneObj.getJSONArray("children");
			loadBones(childrenArr, index);
		}
	}

	public static class Bone {
		public int getParent() {
			return parent;
		}

		public float getLength() {
			return length;
		}

		public void setLength(float length) {
			this.length = length;
		}

		public float getThickness() {
			return thickness;
		}

		public void setThickness(float thickness) {
			this.thickness = thickness;
		}

		public Vector2f getTranslation() {
			return new Vector2f(translation);
		}

		public void setTranslation(Vector2f translation) {
			this.translation.set(translation);
		}

		public float getRotation() {
			return rotation;
		}

		public void setRotation(float rotation) {
			this.rotation = rotation;
		}

		public Vector2f getScale() {
			return new Vector2f(scale);
		}

		public void setScale(Vector2f scale) {
			this.scale.set(scale);
		}

		public int getDepth() {
			return depth;
		}

		public int getChildCount() {
			return children.size();
		}

		public int getChild(int i) {
			return children.get(i);
		}

		public String getName() {
			return name;
		}

		// private constructor - can only be constructed by skeleton
		private Bone(int parent, String name, float length, float thickness, Vector2f translation, float rotation, Vector2f scale, int depth) {
			this.parent = parent;
			this.name = name;
			this.length = length;
			this.thickness = thickness;
			this.translation = new Vector2f(translation);
			this.rotation = rotation;
			this.scale = new Vector2f(scale);
			this.depth = depth;
			children = new ArrayList<Integer>();
		}

		// parent
		private int parent;
		// name of bone
		private String name;
		// length of this bone (for skeleton display purposes)
		private float length;
		// thickness of bone (for skeleton display purposes)
		private float thickness;
		// bind pose translation
		private Vector2f translation;
		// bind pose rotation
		private float rotation;
		// bind pose scale
		private Vector2f scale;
		// depth used to determine drawing order
		private int depth;
		// list of indices of child bones
		private ArrayList<Integer> children;
	}

	////////
	// methods to access parts of skeleton

	public int getBoneCount() {
		return bones.size();
	}

	public Bone getBone(int i) {
		return bones.get(i);
	}

	public int getBoneIndex(String name) {
		Integer index = boneNames.get(name);
		return (index == null) ? -1 : index;
	}

	public Bone getBone(String name) {
		return bones.get(getBoneIndex(name));
	}

	public int getRootBoneCount() {
		return rootBones.size();
	}

	// returns index of the ith root bone
	public int getRootBone(int i) {
		return rootBones.get(i);
	}

	// returns the ith bone which should be drawn
	public int getDrawingOrderedBone(int i) {
		return boneDrawingOrder.get(i);
	}

	////////
	// methods to modify skeleton

	// adds a root bone
	public int addBone(String name, float length, float thickness, Vector2f translation, float rotation, Vector2f scale, int depth) {
		if (boneNames.containsKey(name)) {
			throw new IllegalArgumentException("Bone " + name + " already exists");
		}

		int index = bones.size();

		Bone bone = new Bone(-1, name, length, thickness, translation, rotation, scale, depth);
		bones.add(bone);
		rootBones.add(index);
		addToDrawingOrder(index, depth);
		boneNames.put(name, index);

		return index;
	}

	// adds a child bone
	public int addBone(String name, int parent, float length, float thickness, Vector2f translation, float rotation, Vector2f scale, int depth) {
		if (parent < 0 || parent >= bones.size()) {
			throw new IllegalArgumentException("Invalid parent index");
		}
		if (boneNames.containsKey(name)) {
			throw new IllegalArgumentException("Bone " + name + " already exists");
		}

		int index = bones.size();

		Bone bone = new Bone(parent, name, length, thickness, translation, rotation, scale, depth);
		bones.add(bone);
		bones.get(parent).children.add(index);
		addToDrawingOrder(index, depth);
		boneNames.put(name, index);

		return index;
	}

	// removes a bone
	public void removeBone(int index) {
		if (index < 0 || index >= bones.size()) {
			throw new IllegalArgumentException("Invalid bone index");
		}

		// first remove children by setting to null
		removeBones(index);
		// now find which indices have been removed and remap
		int[] indexMap = new int[bones.size()];
		int next = 0;
		for (int i = 0; i < indexMap.length; ++i) {
			if (bones.get(i) == null) {
				indexMap[i] = -1;
			} else {
				indexMap[i] = next;
				++next;
			}
		}

		// remove null bones
		for (int i = 0; i < bones.size(); ++i) {
			if (bones.get(i) == null) {
				bones.remove(i);
				--i;
			}
		}

		// remap indices
		for (int i = 0; i < rootBones.size(); ++i) {
			int newIndex = indexMap[rootBones.get(i)];
			if (newIndex < 0) {
				rootBones.remove(i);
				--i;
			} else {
				rootBones.set(i, newIndex);
			}
		}

		for (int i = 0; i < boneDrawingOrder.size(); ++i) {
			int newIndex = indexMap[boneDrawingOrder.get(i)];
			if (newIndex < 0) {
				boneDrawingOrder.remove(i);
				--i;
			} else {
				boneDrawingOrder.set(i, newIndex);
			}
		}

		for (Bone b : bones) {
			if (b.parent >= 0) {
				b.parent = indexMap[b.parent];
				assert(b.parent >= 0); // parent should never become invalid
			}
			for (int i = 0; i < b.children.size(); ++i) {
				int newIndex = indexMap[b.children.get(i)];
				if (newIndex < 0) {
					b.children.remove(i);
					--i;
				} else {
					b.children.set(i, newIndex);
				}
			}
		}

		Iterator<Map.Entry<String, Integer>> it = boneNames.entrySet().iterator();
		while (it.hasNext()) {
			Map.Entry<String, Integer> e = it.next();
			int newIndex = indexMap[e.getValue()];
			if (newIndex < 0) {
				it.remove();
			} else {
				e.setValue(newIndex);
			}
		}
	}

	// recursively removes children (sets their array slot to null)
	private void removeBones(int index) {
		Bone bone = bones.get(index);
		bones.set(index, null);
		for (int i = 0; i < bone.children.size(); ++i) {
			removeBones(bone.children.get(i));
		}
	}

	public void setBoneDepth(int index, int depth) {
		if (index < 0 || index >= bones.size()) {
			throw new IllegalArgumentException("Invalid bone index");
		}

		bones.get(index).depth = depth;
		removeFromDrawingOrder(index);
		addToDrawingOrder(index, depth);
	}

	public void setBoneName(int index, String name) {
		if (index < 0 || index >= bones.size()) {
			throw new IllegalArgumentException("Invalid bone index");
		}
		Integer b = boneNames.get(name);
		if (b != null && b != index) {
			throw new IllegalArgumentException("Bone " + name + " already exists");
		}

		Bone bone = bones.get(index);
		boneNames.remove(bone.name);
		bone.name = name;
		boneNames.put(name, index);
	}

	private void addToDrawingOrder(int index, int depth) {
		// add to drawing order list
		int drawIndex;
		for (drawIndex = 0; drawIndex < boneDrawingOrder.size(); ++drawIndex) {
			Bone b = bones.get(boneDrawingOrder.get(drawIndex));
			if (b.depth >= depth) {
				break;
			}
		}
		boneDrawingOrder.add(drawIndex, index);
	}

	private void removeFromDrawingOrder(int index) {
		for (int i = 0; i < boneDrawingOrder.size(); ++i) {
			if (boneDrawingOrder.get(i) == index) {
				boneDrawingOrder.remove(i);
				return;
			}
		}
	}

	// list of bones
	private ArrayList<Bone> bones;
	// list of root bones
	private ArrayList<Integer> rootBones;
	// map of bone names
	private HashMap<String, Integer> boneNames;
	// list of bone indices determining the drawing order
	private ArrayList<Integer> boneDrawingOrder;
}
